ECE 5725 Project :Portable Document Scanner

Work Distribution

Jiaxian Chen

jc3459@cornell.edu

Designed the scanning section.

Bochong Chen

bc446@cornell.edu

Designed the OCR section.

Code Appendix

final.py:

// final.py
import csv
import os
import pickle
import sys
from datetime import datetime

import RPi.GPIO as GPIO
import cv2
import numpy as np
import pygame
import pytesseract

from menu import Menu
from utils import find_biggest, sort_corners, draw_rectangle
from video_stream import VideoStream

# global variables
btn_lock = False
pressed_btn = None
curr_state = Menu.SCAN
menu_lst = [['OCR', 'B & W', 'Take', 'Quit'],
            ['   ', 'Back', 'Save', 'Quit'],
            ['Scan', 'Choose', 'Take', 'Quit'],
            ['Prev', 'Next', 'Apply', 'Quit']]
curr_menu = menu_lst[0]
thresholding = False
take_done = False
main_frame = np.zeros((240, 320, 3))
stream = VideoStream()

# find all the available OCR samples
ocr_filename = []
for filepath in os.listdir("ocr_data"):
    if filepath.endswith('.pkl'):
        ocr_filename.append(filepath.replace(".pkl", ""))
ocr_index = 0
my_data = []

# pre-load the ORB keypoint samples to save time
ocr_sample = cv2.imread(f'ocr_data/{ocr_filename[ocr_index]}.jpg')
ocr_roi = pickle.load(open(f'ocr_data/{ocr_filename[ocr_index]}.pkl', 'rb'))
orb = cv2.ORB_create(1000)
kp1, des1 = orb.detectAndCompute(ocr_sample, None)
bf = cv2.BFMatcher(cv2.NORM_HAMMING)


# convert a OpenCV frame to a pygame compatible one
def cvt2pygame(frame):
    frame = cv2.cvtColor(np.float32(frame), cv2.COLOR_BGR2RGB)
    frame = cv2.resize(frame, (300, 240))  # resize to screen size
    frame = cv2.flip(frame, 1)  # horizontal flip
    frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_COUNTERCLOCKWISE)
    return pygame.surfarray.make_surface(frame)


# interrupt callback function upon button pressed
def button_press_callback(channel):
    global btn_lock, pressed_btn
    if not btn_lock:
        pressed_btn = channel


# change menu state according to the pressed button
def change_state():
    global pressed_btn, curr_state, curr_menu, thresholding, \
        take_done, my_data, ocr_sample, ocr_roi, ocr_index, ocr_filename

    if curr_state == Menu.SCAN:
        if pressed_btn == 22:
            curr_state = Menu.SCAN_TAKE
        elif pressed_btn == 23:
            thresholding = not thresholding
        elif pressed_btn == 27:
            # change back to color mode before going to OCR mode
            thresholding = False
            curr_state = Menu.OCR
            curr_menu = menu_lst[2]
    elif curr_state == Menu.SCAN_TAKE:
        if take_done:  # check if the take action has already finished
            take_done = False
            curr_state = Menu.SCAN_TAKEN
            curr_menu = menu_lst[1]
    elif curr_state == Menu.SCAN_TAKEN:
        if pressed_btn == 22:
            curr_state = Menu.SCAN
            curr_menu = menu_lst[0]
            # save scan results
            time = datetime.now()
            filename = f'saved_data/scanned_document_{str(time.date())}_{time.hour}_{time.minute}_{time.second}.jpg'
            cv2.imwrite(filename, main_frame)
            os.system(f'bypy upload {filename}')
        elif pressed_btn == 23:
            curr_state = Menu.SCAN
            curr_menu = menu_lst[0]
    elif curr_state == Menu.OCR:
        if pressed_btn == 22:
            curr_state = Menu.OCR_TAKE
        elif pressed_btn == 23:
            curr_state = Menu.OCR_CHOOSE
            curr_menu = menu_lst[3]
        elif pressed_btn == 27:
            curr_state = Menu.SCAN
            curr_menu = menu_lst[0]
    elif curr_state == Menu.OCR_TAKE:
        if take_done:  # check if the take action has already finished
            take_done = False
            curr_state = Menu.OCR_TAKEN
            curr_menu = menu_lst[1]
    elif curr_state == Menu.OCR_TAKEN:
        if pressed_btn == 22:
            curr_state = Menu.OCR
            curr_menu = menu_lst[2]
            # save OCR results
            time = datetime.now()
            filename = f'saved_data/ocr_document_{str(time.date())}_{time.hour}_{time.minute}_{time.second}.csv'
            with open(filename, 'w', newline='') as csvfile:
                csv_writer = csv.writer(csvfile)
                csv_writer.writerow([i[3] for i in ocr_roi])
                csv_writer.writerow(my_data)
            os.system(f'bypy upload {filename}')
        elif pressed_btn == 23:
            curr_state = Menu.OCR
            curr_menu = menu_lst[2]
    elif curr_state == Menu.OCR_CHOOSE:
        if pressed_btn == 22:
            curr_state = Menu.OCR
            curr_menu = menu_lst[2]
        elif pressed_btn == 23:
            # switch the OCR sample to the next available
            if ocr_index < len(ocr_filename) - 1:
                ocr_index += 1
                ocr_sample = cv2.imread(f'ocr_data/{ocr_filename[ocr_index]}.jpg')
                ocr_roi = pickle.load(open(f'ocr_data/{ocr_filename[ocr_index]}.pkl', 'rb'))
        elif pressed_btn == 27:
            # switch the OCR sample to the previous available
            if ocr_index > 0:
                ocr_index -= 1
                ocr_sample = cv2.imread(f'ocr_data/{ocr_filename[ocr_index]}.jpg')
                ocr_roi = pickle.load(open(f'ocr_data/{ocr_filename[ocr_index]}.pkl', 'rb'))

    # state change is finished, erase pressed_btn
    pressed_btn = None


# obtain frame depending on menu state
def get_frame():
    global main_frame, curr_state

    if curr_state == Menu.SCAN or curr_state == Menu.OCR:
        main_frame, success = get_scanning_frame(False)
    elif curr_state == Menu.SCAN_TAKE:
        main_frame, success = get_scanning_frame(True)
    elif curr_state == Menu.OCR_TAKE:
        main_frame, myData = get_ocr_frame()
    elif curr_state == Menu.SCAN_TAKEN or curr_state == Menu.OCR_TAKEN:
        pass
    elif curr_state == Menu.OCR_CHOOSE:
        main_frame = get_ocr_sample_frame()


# load the sample picture and overlay the fields
def get_ocr_sample_frame():
    img_show = ocr_sample.copy()
    img_mask = np.zeros_like(img_show)

    for x, r in enumerate(ocr_roi):
        cv2.rectangle(img_mask, r[0], r[1], (0, 255, 0), cv2.FILLED)
        img_show = cv2.addWeighted(img_show, 0.99, img_mask, 0.1, 0)  # overlay ROI rectangles
        cv2.putText(img_show, r[3], r[0], cv2.FONT_HERSHEY_PLAIN, 3, (0, 0, 255), 3)  # overlay field text

    return img_show


# acquire a scanning preview or result frame
def get_scanning_frame(get_result=False):
    global stream, thresholding, take_done

    img = stream.read().copy()
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # turn gray for later edge detection
    img_blur = cv2.GaussianBlur(img_gray, (5, 5), 1)  # blur to remove high frequency noises
    img_threshold = cv2.Canny(img_blur, 100, 150)  # get picture with edges only
    kernel = np.ones((5, 5))
    img_dial = cv2.dilate(img_threshold, kernel, iterations=2)
    img_threshold = cv2.erode(img_dial, kernel, iterations=1)

    contours, _ = cv2.findContours(img_threshold, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # find closed contours
    biggest, maxArea = find_biggest(contours)  # find the biggest contour with 4 sides

    if get_result and biggest.size != 0:  # this means the program wants the scanned result
        biggest = sort_corners(biggest)
        pts1 = np.float32(biggest)
        pts2 = np.float32(
            [[0, 0], [img.shape[1], 0], [0, img.shape[0]], [img.shape[1], img.shape[0]]])
        matrix = cv2.getPerspectiveTransform(pts1, pts2)
        img = cv2.warpPerspective(img, matrix, (img.shape[1], img.shape[0]))  # perform perspective transform

        if thresholding:
            # apply adaptive thresholding
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            img = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)

        take_done = True
        return img, True

    else:  # this means the program just want a preview frame
        if thresholding:
            # apply adaptive thresholding
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            img = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

        if biggest.size == 0:
            cv2.drawContours(img, contours, -1, (0, 255, 0), 5)
        else:
            biggest = sort_corners(biggest)
            img = draw_rectangle(img, biggest, 5)

        return img, False


def get_ocr_frame():
    global stream, take_done, ocr_sample, ocr_roi, orb, kp1, des1, bf, my_data

    # ORB is still not robust enough
    # for now, perform document detection first, then ORB
    img, success = get_scanning_frame(True)
    while not success:
        img, success = get_scanning_frame(True)

    img = stream.read().copy()

    kp2, des2 = orb.detectAndCompute(img, None)

    matches = bf.match(des2, des1)
    matches.sort(key=lambda x: x.distance)
    good_matches = matches[:int(len(matches) * 0.25)]

    src_points = np.float32([kp2[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
    dest_points = np.float32([kp1[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
    M, _ = cv2.findHomography(src_points, dest_points, cv2.RANSAC, 5.0)
    img_scan = cv2.warpPerspective(img, M, (ocr_sample.shape[1], ocr_sample.shape[0]))

    img_scan = cv2.resize(img_scan, (ocr_sample.shape[1], ocr_sample.shape[0]))
    img_show = img_scan.copy()
    img_mask = np.zeros_like(img_show)

    my_data = []

    for x, r in enumerate(ocr_roi):
        cv2.rectangle(img_mask, r[0], r[1], (0, 255, 0), cv2.FILLED)
        img_show = cv2.addWeighted(img_show, 0.99, img_mask, 0.1, 0)

        img_crop = img_scan[r[0][1]: r[1][1], r[0][0]: r[1][0]]

        if r[2] == 'text':
            my_data.append(pytesseract.image_to_string(img_crop)
                           .replace("\x0c", "").replace("\n", ""))
        elif r[2] == 'box':
            imgGray = cv2.cvtColor(img_crop, cv2.COLOR_BGR2GRAY)
            imgThreshold = cv2.threshold(imgGray, 170, 255, cv2.THRESH_BINARY_INV)[1]
            isChecked = cv2.countNonZero(imgThreshold) > 0.5 * imgGray.shape[0] * imgGray.shape[1]
            my_data.append(isChecked)
        cv2.putText(img_show, str(my_data[x]), r[0], cv2.FONT_HERSHEY_PLAIN, 3, (0, 0, 255), 3)
    print(my_data)
    take_done = True
    return img_show, my_data


# main loop
def main():
    global btn_lock, pressed_btn, curr_state, stream, main_frame, curr_menu

    # os.putenv('SDL_VIDEODRIVER', 'fbcon')  # Display on piTFT
    # os.putenv('SDL_FBDEV', '/dev/fb1')  #
    # os.putenv('SDL_MOUSEDRV', 'TSLIB')  # Track mouse clicks on piTFT
    # os.putenv('SDL_MOUSEDEV', '/dev/input/touchscreen')

    btn_pins = [17, 22, 23, 27]
    GPIO.setmode(GPIO.BCM)
    GPIO.setup(btn_pins, GPIO.IN, pull_up_down=GPIO.PUD_UP)

    WHITE = (255, 255, 255)
    BLACK = (0, 0, 0)
    text_pos = [(305, 200), (305, 140), (305, 80), (305, 20)]

    pygame.init()
    pygame.mouse.set_visible(False)
    screen = pygame.display.set_mode((320, 240))
    screen.fill(BLACK)
    font = pygame.font.Font(None, 20)

    for btn_pin in btn_pins:  # bind button press events
        GPIO.add_event_detect(btn_pin, GPIO.FALLING, callback=button_press_callback, bouncetime=300)

    stream.start()

    pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"

    try:
        while True:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    sys.exit()

            btn_lock = True
            if pressed_btn == 17:
                sys.exit()
            change_state()
            btn_lock = False

            get_frame()

            screen.fill(BLACK)
            screen.blit(cvt2pygame(main_frame), (0, 0))
            for i in range(4):
                text = font.render(curr_menu[i], True, WHITE)
                text = pygame.transform.rotate(text, 90)
                screen.blit(text, text_pos[i])

            pygame.display.flip()

    finally:
        print("Program finished")
        stream.cleanup()
        cv2.destroyAllWindows()
        pygame.quit()


if __name__ == "__main__":
    main()

              

menu.py:

from enum import Enum, unique
@unique
class Menu(Enum):
    SCAN = 1
    SCAN_TAKE = 2
    SCAN_TAKEN = 3
    OCR = 4
    OCR_TAKE = 5
    OCR_TAKEN = 6
    OCR_CHOOSE = 7

                

utils:

import cv2
import numpy as np


def sort_corners(myPoints):
    myPoints = myPoints.reshape((4, 2))
    myPointsNew = np.zeros((4, 1, 2), dtype=np.int32)
    add = myPoints.sum(1)

    myPointsNew[0] = myPoints[np.argmin(add)]
    myPointsNew[3] = myPoints[np.argmax(add)]
    diff = np.diff(myPoints, axis=1)
    myPointsNew[1] = myPoints[np.argmin(diff)]
    myPointsNew[2] = myPoints[np.argmax(diff)]

    return myPointsNew


def find_biggest(contours):
    biggest = np.array([])
    max_area = 0
    for i in contours:
        area = cv2.contourArea(i)
        if area > 5000:
            peri = cv2.arcLength(i, True)
            approx = cv2.approxPolyDP(i, 0.02 * peri, True)
            if area > max_area and len(approx) == 4:
                biggest = approx
                max_area = area
    return biggest, max_area


def draw_rectangle(img, biggest, thickness):
    cv2.line(img, (biggest[0][0][0], biggest[0][0][1]), (biggest[1][0][0], biggest[1][0][1]), (0, 255, 0), thickness)
    cv2.line(img, (biggest[0][0][0], biggest[0][0][1]), (biggest[2][0][0], biggest[2][0][1]), (0, 255, 0), thickness)
    cv2.line(img, (biggest[3][0][0], biggest[3][0][1]), (biggest[2][0][0], biggest[2][0][1]), (0, 255, 0), thickness)
    cv2.line(img, (biggest[3][0][0], biggest[3][0][1]), (biggest[1][0][0], biggest[1][0][1]), (0, 255, 0), thickness)

    return img

videostream.py:

                import cv2.cv2
import numpy as np


class VideoStream:
    def __init__(self):
        self.cap = cv2.VideoCapture(0)
        self.cap.set(cv2.cv2.CAP_PROP_FRAME_WIDTH, 1024)
        self.cap.set(cv2.cv2.CAP_PROP_FRAME_HEIGHT, 576)
        self.stopped = False
        self.frame = np.zeros((576, 1024, 3), np.uint8)

    def start(self):
        t = Thread(target=self.update, args=())
        t.daemon = True
        t.start()
        return self

    def update(self):
        while True:
            if self.stopped:
                return

            ret, new_frame = self.cap.read()
            if not ret:
                self.stop()
                return
            self.frame = new_frame

    def read(self):
        return self.frame

    def stop(self):
        self.stopped = True

    def cleanup(self):
        self.cap.release()